#include <cuda_runtime.h>
#include <iostream>

__global__ void AddData1toData0(float* data0, float data1) {
    // 先读取原始值
    float originalValue = *data0;

    // 打印原始值（仅用于调试）
    printf("Original value on GPI0/GPU1: %f\n", originalValue);
    printf("add data1: %f\n", data1);

    // 再次读取修改后的值
    *data0 = originalValue + data1;

    // 打印最终值（仅用于调试）
    printf("Final value on GPI0/GPU1: %f\n", *data0);
}

int main() {
    int numDevices;
    cudaGetDeviceCount(&numDevices);
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    if (numDevices < 2) {
        std::cerr << "This demo requires at least two GPUs." << std::endl;
        return -1;
    }

    const int device0 = 0; // 第一个设备
    const int device1 = 1; // 第二个设备

    // 检查设备是否支持 P2P 访问
    int canAccessPeer;
    cudaSetDevice(device0);
    cudaDeviceCanAccessPeer(&canAccessPeer, device0, device1);
    if (!canAccessPeer) {
        std::cerr << "Device " << device0 << " cannot access peer memory of Device " << device1 << "." << std::endl;
        return -1;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    // 启用 P2P 访问
    cudaSetDevice(device0);
    cudaDeviceEnablePeerAccess(device1, 0);

    cudaSetDevice(device1);
    cudaDeviceEnablePeerAccess(device0, 0);

    // 在 GPU0 上分配内存
    float* dev0Data;
    cudaSetDevice(device0);
    cudaMalloc(&dev0Data, sizeof(float));

    // 初始化 GPU0 上的数据
    float hostData0 = 10.0f; // 初始值设置为 10.0
    float hostData1 = 20.0f; // 初始值设置为 20.0
    cudaMemcpy(dev0Data, &hostData0, sizeof(float), cudaMemcpyHostToDevice);
    printf("[%s +%d %s] Host-device0 Value=%f\n", __FILE__, __LINE__, __func__, hostData0);

    // 在 GPU1 上启动 Kernel 修改 GPU0 的数据
    cudaSetDevice(device1);
    printf("[%s +%d %s] before AddData1toData0 Host-device1 Value=%f GPU1 add to GPU0 and thre result=%f\n", __FILE__, __LINE__, __func__, hostData1, hostData0 + hostData1);
    AddData1toData0<<<1, 1>>>(dev0Data, hostData1);

    // 同步以确保 Kernel 执行完成
    cudaSetDevice(device1);
    cudaDeviceSynchronize();

    // 将 GPU0 的数据复制回主机
    float result;
    cudaSetDevice(device0);
    cudaMemcpy(&result, dev0Data, sizeof(float), cudaMemcpyDeviceToHost);
    printf("[%s +%d %s] after kernel-calc and memcpy form Device0 Host-device0 result=%f\n", __FILE__, __LINE__, __func__, result);

    // 输出结果
    std::cout << "Final value on GPU0: " << result << std::endl;

    // 清理资源
    cudaSetDevice(device0);
    cudaFree(dev0Data);
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    return 0;
}

